import torch
import torch.nn as nn
from mmcls.models import BACKBONES
from mmcls.models.backbones import VisionTransformer
from mmcls.models.utils import resize_pos_embed
from typing import List


device = 'cuda:0'

@BACKBONES.register_module()
class IMT(VisionTransformer):

    def __init__(self,
                 prompt_length: int = 1,
                 prompt_layers: List[int] = None,
                 prompt_pos: str = 'postpend',
                 prompt_init: str = 'token',
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        for param in self.parameters():
            param.requires_grad = False


        self.prompt_layers = [0] if prompt_layers is None else prompt_layers  


        prompt = torch.empty(len(self.prompt_layers), 577, prompt_length)


        if prompt_init == 'uniform':
            nn.init.uniform_(prompt, -0.08, 0.08)
        elif prompt_init == 'zero':
            nn.init.zeros_(prompt)
        elif prompt_init == 'kaiming':
            nn.init.kaiming_normal_(prompt)
        elif prompt_init == 'token':
            nn.init.zeros_(prompt)
            self.prompt_initialized = False
        else:
            nn.init.normal_(prompt, std=0.02)
        self.prompt = nn.Parameter(prompt, requires_grad=True)
        self.prompt_length = prompt_length
        self.prompt_pos = prompt_pos


    def forward(self, x):
        """Following mmcls implementation."""
        B = x.shape[0] # batch size
        x, patch_resolution = self.patch_embed(x)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1) # x: (B,577,768)
        x = x + resize_pos_embed(
            self.pos_embed,
            self.patch_resolution,
            patch_resolution,
            mode=self.interpolate_mode,
            num_extra_tokens=self.num_extra_tokens) # x: (B,577,768)
        x = self.drop_after_pos(x) # x: (B,577,768)




        outs = []
        for i, layer in enumerate(self.layers):
            if i in self.prompt_layers:

                def IMT_attention(layer_attn, x):
                    B, N, _ = x.shape # x:(4,577,768)
                    x = layer.norm1(x)
                    # Add prompt
                    if hasattr(self, 'prompt_initialized') and not self.prompt_initialized:
                        with torch.no_grad():
                            self.prompt.data += x.mean([0, 2]).detach().clone()[None, :, None]
                        self.prompt_initialized = True
                    # prompt: (1=layer,4,577,1)
                    prompt = self.prompt.unsqueeze(1).expand(-1, x.shape[0], -1, -1)
                    # Add prompt
                    prompt = prompt[i, :, :, :].clone().expand(B, N, self.embed_dims)
                    qkv_ = layer_attn.qkv(prompt) # (4,577,1) -> qkv:(4,577,2304)->(4,577,3)


                    qkv = layer_attn.qkv(x) # qkv:(4,577,2304)
                    qkv = qkv+qkv_
                    qkv = qkv.reshape(B, N, 3, layer_attn.num_heads,layer_attn.head_dims).permute(2, 0, 3, 1, 4) # qkv:(3,4,12,577,64)
                    q, k, v = qkv[0], qkv[1], qkv[2] #(4,12,577,64)

                    attn = (q @ k.transpose(-2, -1)) * layer_attn.scale
                    attn = attn.softmax(dim=-1)
                    attn = layer_attn.attn_drop(attn)

                    x = (attn @ v).transpose(1, 2).reshape(B, N, layer_attn.embed_dims)
                    x = layer_attn.proj(x)
                    x = layer_attn.out_drop(layer_attn.proj_drop(x))

                    if layer_attn.v_shortcut:
                        x = v.squeeze(1) + x
                    return x

                # x = x + layer.attn(layer.norm1(x))
                x = x + IMT_attention(layer.attn, x)
                x = layer.ffn(layer.norm2(x), identity=x)
            else:    
                x = layer(x)

            if i == len(self.layers) - 1 and self.final_norm:
                x = self.norm1(x)

            if i in self.out_indices:
                outs.append(x[:, 0])

        return tuple(outs)